#!/usr/bin/python

# avl2postgis.py: load AVL data (currently NextBus) into PostGIS
# Copyright 2011 Matt Conway

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#   http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Authors:
# Matt Conway: main code

# The table you connect to should look like this:
# agency varchar
# vehicle varchar
# direction varchar
# heading int
# route varchar
# time timestamp
# the_geom

from optparse import OptionParser
from pyquery import PyQuery as pq
from datetime import datetime
from time import sleep
from psycopg2 import connect
from urllib2 import urlopen
import sys

# Parse the options
parser = OptionParser()
parser.add_option('-a', '--agency', dest='agency', type='string',
                  help='the NextBus agency name (e.g. sf-muni)',
                  metavar='AGENCY', default='sf-muni')
parser.add_option('-t', '--table', dest='table', type='string',
                  help='the table to store data in',
                  metavar='TABLE', default='nextbus')
parser.add_option('-d', '--dsn', dest='dsn', type='string',
                  help='database dsn, as documented at http://initd.org/psycopg/docs/module.html#psycopg2.connect',
                  metavar='DSN', default='dbname=arrivals')
parser.add_option('-i', '--interval', dest='interval', type='int',
                  help='number of seconds to pause between NextBus requests',
                  metavar='interval', default=30)
parser.add_option('--all-vehicles', dest='allvehicles',
                  help='query all vehicles each time, instead of just updated ones',
                  default=False)
parser.add_option('--no-deltas', dest='nodeltas',
                  help="Don't use deltas to calculate time, just put in retrieval time",
                  default=False)

options, args = parser.parse_args()

# start with all reported locations in last 15 mins
since_time = 0

# This function grabs the data from NextBus
def getDataFromNextBus():
    global since_time, options

    if options.allvehicles:
        qtime = 0
    else:
        qtime = since_time

    # Build the NB URL
    nburl = "http://webservices.nextbus.com/service/publicXMLFeed?command=vehicleLocations&a={agency}&t={since_time}".format(agency=options.agency, since_time=qtime)

    # Get the XML
    x = pq(urlopen(nburl).read())

    since_time = int(x('lastTime').attr('time'))
    
    def makeDBObj(vRaw):
        # Create a new pquery obj, we just get passed a DOM obj
        v = pq(vRaw)

        # Figure out the time
        # This is the epoch time this location was recorded
        if (options.nodeltas == False):
            epochReport = since_time - (int(v.attr('secsSinceReport')) * 1000)
        else:
            epochReport = since_time

        # Get in seconds, which is what python likes
        utctime = epochReport / 1000
        
        # build a datetime obj
        dt = datetime.utcfromtimestamp(utctime)

        return dict(
            agency=options.agency,
            vehicle=v.attr('id'),
            direction=v.attr('dirTag'),
            heading=int(v.attr('heading')),
            route=v.attr('routeTag'),
            # zfill is so that month is 05, not 5
            time='{year}-{month}-{day} {hour}:{minute}:{second}'.format(  
                year=dt.year,
                month=str(dt.month).zfill(2),
                day=str(dt.day).zfill(2),
                hour=str(dt.hour).zfill(2),
                minute=str(dt.minute).zfill(2),
                second=str(dt.second).zfill(2),
                ),

            # Note that it needs to be in EWKT format, so it can be put into
            # PostGIS
            wktgeom='SRID=4326;POINT({lon} {lat})'.format(lon=v.attr('lon'),
                                                 lat=v.attr('lat'))
            )

    return map(makeDBObj, x('vehicle'))

def addVehicleToPostGIS(v):
    # change to print cursor.mogrify to see what SQL will be executed
    # BTW, don't let anyone type in a table name; that could create the
    # possibility of an SQL injection attack!
    cursor.execute("""
INSERT INTO {table} (agency, vehicle, direction, heading, route, 
   time, the_geom)
VALUES (%(agency)s, %(vehicle)s, %(direction)s, %(heading)s, %(route)s, %(time)s,
   ST_GeomFromEWKT(%(wktgeom)s));
""".format(table=options.table), dict(table=options.table, **v))

# Connect to the DB
conn = connect(options.dsn)
cursor = conn.cursor()

# Loop forever, or until killed
while 1:
    try:
        g = getDataFromNextBus()
        print 'Adding {l} points to database'.format(l=len(g))
        for v in g:
            addVehicleToPostGIS(v)
        # Commit
        conn.commit()
    except:
        print sys.exc_info()
        print 'Exception occured - skipping this iteration (connection reset by peer?)'

    # The reason this is inside a try...except is b/c
    # sometimes people want to skip the sleep and go straight to the
    # next iteration.
    try:
        sleep(options.interval)
    except:
        pass
        

